import sys
import piexif
from PIL import Image
import io

def embed_image_as_exif_thumbnail(input_thumbnail_path, input_original_path, output_path, max_thumbnail_size_kb=64):
    """
    使用 piexif 将一张图片嵌入到另一张 JPEG 图片的 EXIF 数据中。
    如果图片过大，会先进行压缩。
    
    Args:
        input_thumbnail_path: 用作嵌入图片的文件路径。
        input_original_path: 原始 JPEG 图片文件路径，其 EXIF 数据将被修改。
        output_path: 修改后带有新 EXIF 数据的图片保存路径。
        max_thumbnail_size_kb: EXIF缩略图的最大允许大小（KB）。
    """
    max_size_bytes = max_thumbnail_size_kb * 1024  # 转换为字节

    # 1. 读取用作嵌入的图片，并转换为符合大小要求的 JPEG 字节数据
    try:
        with Image.open(input_thumbnail_path) as img:
            img = img.convert("RGB")  # 确保为RGB模式
            thumb_buffer = io.BytesIO()
            
            # 初始质量
            quality = 95
            # 先尝试保存，看看大小
            img.save(thumb_buffer, format="JPEG", quality=quality)
            thumbnail_data = thumb_buffer.getvalue()

            # 如果图片太大，就降低质量并重试
            while len(thumbnail_data) > max_size_bytes and quality > 20:
                thumb_buffer = io.BytesIO()
                quality -= 5  # 每次降低5个质量点
                img.save(thumb_buffer, format="JPEG", quality=quality)
                thumbnail_data = thumb_buffer.getvalue()

            if len(thumbnail_data) > max_size_bytes:
                print(f"警告: 即使将质量降至 {quality}，嵌入图片仍大于 {max_thumbnail_size_kb}kB。")
                print(f"原始大小: {len(thumbnail_data)} 字节, 最大大小: {max_size_bytes} 字节")
                # 如果还是太大，可以考虑调整图片尺寸
                # 例如，将图片缩小到最大边为 500 像素
                max_dimension = 500
                img.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
                thumb_buffer = io.BytesIO()
                img.save(thumb_buffer, format="JPEG", quality=quality)
                thumbnail_data = thumb_buffer.getvalue()

            if len(thumbnail_data) > max_size_bytes:
                print(f"错误: 无法将嵌入图片压缩至 {max_thumbnail_size_kb}kB 以内。")
                sys.exit(1)

            print(f"嵌入图片已压缩至 {len(thumbnail_data)} 字节 (质量: {quality})")
    except Exception as e:
        print(f"错误: 无法读取或处理嵌入图片 '{input_thumbnail_path}': {e}")
        sys.exit(1)

    # 2. 读取原始图片的 EXIF 数据
    try:
        with Image.open(input_original_path) as img:
            exif_dict = piexif.load(img.info.get("exif", b""))
    except Exception as e:
        print(f"错误: 无法读取原始图片 '{input_original_path}': {e}")
        sys.exit(1)

    # 3. 将压缩后的图片数据放入 EXIF 的 thumbnail 字段
    exif_dict["thumbnail"] = thumbnail_data

    # 4. 将修改后的 EXIF 数据重新插入到图片中并保存
    try:
        exif_bytes = piexif.dump(exif_dict)
        with Image.open(input_original_path) as img:
            if img.mode != "RGB":
                img = img.convert("RGB")
            img.save(output_path, "jpeg", exif=exif_bytes, quality=95, optimize=True)
        print(f"成功: 已将 '{input_thumbnail_path}' 嵌入到 '{input_original_path}' 的 EXIF 中，输出为 '{output_path}'")
    except Exception as e:
        print(f"错误: 无法保存输出图片 '{output_path}': {e}")
        sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("用法: python script.py <输入嵌入图片路径> <输入原图路径> <输出文件路径>")
        sys.exit(1)

    input_embed = sys.argv[1]
    input_orig = sys.argv[2]
    output_file = sys.argv[3]

    embed_image_as_exif_thumbnail(input_embed, input_orig, output_file)